package com.www.ve.utils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.BiConsumer;
import java.util.function.Function;

/**
 * 树形结构工具类
 *
 * @author XiaoJie
 * @date 2025/5/23 15:09
 */
public class TreeBuilderUtils {

    /**
     * 将数据转换为树状数据
     * @param list  原始数据
     * @param superGetter   上级标识字段（例：User::getParentId）
     * @param thisGetter    当前标识字段（例：User::getId）
     * @param nextSetter    下级数据存储字段（例：User::setUsers）
     * @return  树形数据
     */
    public static <T, R> List<T> tree(List<T> list, Function<T, R> superGetter, Function<T, R> thisGetter, BiConsumer<T, List<T>> nextSetter) {
        Map<R, List<T>> map = groupingBy(list, superGetter);
        List<T> tmp = new ArrayList<>(list);
        for (T node : list) {
            R apply = thisGetter.apply(node);
            List<T> ts = map.get(apply);
            if (ts != null && ts.size() > 0) {
                nextSetter.accept(node, ts);
                tmp.removeIf(ts::contains);
            }
        }
        return tmp;
    }

    /**
     * 将数据转换为树状数据
     * @param list  原始数据
     * @param root  根数据(顶级节点)
     * @param superGetter   上级标识字段（例：User::getParentId）
     * @param thisGetter    当前标识字段（例：User::getId）
     * @param nextSetter    下级数据存储字段（例：User::setUsers）
     * @return  树形数据
     */
    public static <T, R> List<T> tree(List<T> list, R root, Function<T, R> superGetter, Function<T, R> thisGetter, BiConsumer<T, List<T>> nextSetter) {
        Map<R, List<T>> map = groupingBy(list, superGetter);
        List<T> result = new ArrayList<>();
        for (T node : list) {
            R apply = thisGetter.apply(node);
            List<T> ts = map.get(apply);
            if (ts != null && ts.size() > 0) {
                nextSetter.accept(node, ts);
            }
            R superValue = superGetter.apply(node);
            if (Objects.deepEquals(superValue, root)) {
                result.add(node);
            }
        }
        return result;
    }

    private static <T, R> Map<R, List<T>> groupingBy(List<T> list, Function<T, R> superGetter) {
        Map<R, List<T>> map = new HashMap<>();
        if(list == null) {
            return map;
        }
        for (T t : list) {
            R apply = superGetter.apply(t);
            List<T> ts = map.computeIfAbsent(apply, k -> new ArrayList<>());
            ts.add(t);
        }
        return map;
    }
}
